package SecurityMessage;

import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.PublicKey;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;

/**
 * @创建人 MBY381
 * @创建时间 2022/6/23 15:33
 * @描述 用于存储此系统的公钥库（RSA公钥）以及端口库（同机不同端口不同套接字通信）
 */

public class KeyServer {
    public static Map<Integer, PublicKey> publicKeyMap;
    private static Map<Integer, String> macKeyStringMap;
    public static HashSet<Integer> portLib;
    private static ServerSocket serverSocket;

    public static void main(String[] args) throws IOException {
        publicKeyMap = new HashMap<>();
        portLib = new HashSet<>();
        serverSocket = new ServerSocket(8888);
        macKeyStringMap = new HashMap<>();
        receive();
    }

    public static void receive() throws IOException {
        while (true) {
            Socket socket = serverSocket.accept();
            InputStream inputStream = socket.getInputStream();
            OutputStream outputStream = socket.getOutputStream();

            try {
                byte[] bytes = new byte[256];
                int len = inputStream.read(bytes);
                String tag = new String(bytes, 0, len);
                System.out.println("服务器接收到一个操作请求：" + tag);
                switch (tag) {
                    case "updateKey":
                        returnSuccess(outputStream);
                        updatePubKey(inputStream);
                        break;
                    case "login":
                        returnSuccess(outputStream);
                        login(inputStream);
                        returnSuccess(outputStream);
                        break;
                    case "logout":
                        returnSuccess(outputStream);
                        logout(inputStream);
                        returnSuccess(outputStream);
                        break;
                    case "getKey":
                        returnSuccess(outputStream);
                        getKey(inputStream);
                        break;
                    case "checkOnline":
                        returnSuccess(outputStream);
                        checkOnline(inputStream, outputStream);
                        break;
                    case "getMacKey":
                        returnSuccess(outputStream);
                        getMacKey(inputStream, outputStream);
                        break;
                    case "uploadMacKey":
                        returnSuccess(outputStream);
                        uploadMacKey(inputStream, outputStream);
                        break;

                }
                System.out.println("目前在线用户:" + portLib);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private static void uploadMacKey(InputStream inputStream, OutputStream outputStream) throws IOException {
        Integer keyPort = Integer.parseInt(receiveMessage(inputStream));
        returnSuccess(outputStream);
        try {
            byte[] keyBytes = new byte[256];
            int len = inputStream.read(keyBytes);
            String macKey = new String(keyBytes, 0, len).trim();

            System.out.println(keyPort + " 更新了它的MAC密钥（已经RSA加密） " + macKey);
            macKeyStringMap.put(keyPort, macKey);
            returnSuccess(outputStream);
            System.out.println();
        } catch (IOException e) {
            returnFalse(outputStream);
            e.printStackTrace();
        }
    }

    private static void getMacKey(InputStream inputStream1, OutputStream outputStream1) throws IOException {
        Integer keyPort = Integer.parseInt(receiveMessage(inputStream1));
        if (macKeyStringMap.containsKey(keyPort)) {
            returnSuccess(outputStream1);
        } else {
            returnFalse(outputStream1);
            return;
        }

        Socket socket = serverSocket.accept();
        OutputStream outputStream = socket.getOutputStream();
        try {
            outputStream.write(macKeyStringMap.get(keyPort).getBytes());
            System.out.println("\nServer发送MAC密钥" + macKeyStringMap.get(keyPort) + "\n");
        } catch (IOException e) {
            returnFalse(outputStream);
            e.printStackTrace();
        }
    }

    //公钥库更新
    public static void updatePubKey(InputStream inputStream1) throws IOException {

        Integer keyPort = Integer.parseInt(receiveMessage(inputStream1));

        Socket socket = serverSocket.accept();
        OutputStream outputStream = socket.getOutputStream();
        System.out.println("开始接收公钥更新\n");

        try {
            InputStream inputStream = socket.getInputStream();
            ObjectInputStream objectInputStream = new ObjectInputStream(inputStream);

            PublicKey publicKey = (PublicKey) objectInputStream.readObject();
            System.out.println(keyPort + "更新了它的公钥" + publicKey.getAlgorithm());
            publicKeyMap.put(keyPort, publicKey);
            System.out.println("目前公钥服务器中有 " + publicKeyMap.size() + " 个公钥\n" + "分布如下" + publicKeyMap.keySet());
            returnSuccess(outputStream);
            System.out.println();
        } catch (IOException | ClassNotFoundException e) {
            returnFalse(outputStream);
            e.printStackTrace();
        }
    }

    public static void getKey(InputStream inputStream1) throws IOException {
        System.out.println("\n收到一个公钥查询请求 -->\n");
        //先接收port
        Integer targetPort = Integer.parseInt(receiveMessage(inputStream1));
        System.out.println("已接收目标公钥ID： " + targetPort);
        inputStream1.close();

        Socket socket = serverSocket.accept();
        OutputStream outputStream = socket.getOutputStream();
        try {
            InputStream inputStream = socket.getInputStream();

            ObjectOutputStream objectOutputStream = new ObjectOutputStream(outputStream);
            objectOutputStream.writeObject(publicKeyMap.get(targetPort));

        } catch (IOException e) {
            returnFalse(outputStream);
            e.printStackTrace();
        }
    }


    public static void checkOnline(InputStream inputStream, OutputStream outputStream) throws IOException {

        byte[] bytes = new byte[1024];
        int len = inputStream.read(bytes);
        Integer targetPort = Integer.valueOf(new String(bytes, 0, len));

        if (portLib.contains(targetPort)) {
            returnSuccess(outputStream);
        } else {
            returnFalse(outputStream);
        }
    }


    public static void login(InputStream inputStream) throws IOException {
        Integer integer = Integer.parseInt(receiveMessage(inputStream));
        portLib.add(integer);
        System.out.println("用户 " + integer + " 登录成功\n");
    }


    public static void logout(InputStream inputStream) throws IOException {
        Integer integer = Integer.parseInt(receiveMessage(inputStream));
        portLib.remove(integer);
        System.out.println("用户 " + integer + " 登出成功\n");
    }

    public static String receiveMessage(InputStream inputStream) throws IOException {
        byte[] bytes1 = new byte[256];
        int length = inputStream.read(bytes1);
        return new String(bytes1, 0, length).trim();
    }

    public static void returnSuccess(OutputStream outputStream) {
        try {
            outputStream.write("success".getBytes(StandardCharsets.UTF_8));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void returnFalse(OutputStream outputStream) throws IOException {
        outputStream.write("false".getBytes(StandardCharsets.UTF_8));
    }
}
